# -*- coding: utf-8 -*-
# @Author  : longbhu
# @Time    : 2025/2/24 13:06
# @Function:
import rasterio
import json
import os
import numpy as np
import pandas as pd

from typing import List, Union, Optional
from rasterio import float32
from calc_v1 import Parameter
from typing import Dict

# 加载 JSON 数据
def load_parameters_from_json(file_path: str) -> Dict:
    """
    从指定路径的 JSON 文件中加载参数。

    参数:
    file_path (str): JSON 文件的路径。

    返回:
    Dict: 包含参数数据的字典。
    """
    try:
        with open(file_path, 'r', encoding='utf-8') as file:
            data = json.load(file)
    except UnicodeDecodeError:
        print(f"尝试使用 UTF-8 编码读取文件失败，尝试使用 GBK 编码读取...")
        with open(file_path, 'r', encoding='gbk') as file:
            data = json.load(file)
    return data


def load_parameters(json_file_path: str, mode: str) -> List[Parameter]:
    """
    从 JSON 文件中加载参数并创建 Parameter 对象列表。

    参数:
    json_file_path (str): 包含参数的 JSON 文件路径。

    返回:
    List[Parameter]: Parameter 对象列表。
    """
    json_data = load_parameters_from_json(json_file_path)
    parameters = []

    for var in json_data[mode]:
        var_name = var['name']
        if var['is_matrix'] and 'file_path' in var:
            try:
                with rasterio.open(var['file_path']) as src:
                    value = src.read(1)  # 假设 GeoTIFF 文件只有一个波段
            except Exception as e:
                print(f"{var_name}: 加载文件 {var['file_path']} 时出错: {e}")
                value = None
        elif not var['is_matrix'] and 'file_path' in var:
            file_path = var['file_path']
            if file_path.endswith(('.xlsx', '.xls')):
                try:
                    df = pd.read_excel(file_path)
                    value = dict(zip(df['code'], df['value']))
                    print(f"{var_name}: 将 Excel 数据读取成字典: {value}")
                except Exception as e:
                    print(f"{var_name}: 加载文件 {file_path} 时出错: {e}")
                    value = None
            else:
                value = var['value']
        else:
            value = var['value']

        # 创建 Parameter 对象并添加到参数列表
        parameter = Parameter(
            name=var['name'],
            value=value,
            is_matrix=var['is_matrix'],
            used_in_formula=var['used_in_formula'],
            description=var['description'],
            is_excel=var['is_excel'],
            is_value_parameter=var['is_value_parameter']
        )
        parameters.append(parameter)

    return parameters


def calculate_with_excel_mapping(ecosystem_array: np.ndarray, excel_dict: dict, formula_result: Union[float, np.ndarray]) -> np.ndarray:
    """
    使用Excel字典映射替换生态图中的值，并乘以公式结果。

    :param ecosystem_array: 生态图数组
    :param excel_dict: Excel字典映射
    :param formula_result: 公式计算结果
    :return: 修改后的生态图数组
    """
    # 创建原始数组的副本，避免修改原矩阵的值
    modified_array = ecosystem_array.copy().astype(np.float32)

    # 根据Excel字典映射替换副本中的值
    for key, value in excel_dict.items():
        if value != -9999:  # 跳过 value 为 -9999 的情况
            print(f"将 {key} 替换为 {value}")
            modified_array[modified_array == key] = value
        else:
            # todo 效率太低了，需要优化
            value = 0
            print(f"将 {key} 替换为 {value}")
            modified_array[modified_array == key] = value

            # masked_data = np.where(mask, data, 0).astype(data.dtype)

    # 将修改后的数组乘以公式结果
    if isinstance(formula_result, (float, int)):
        modified_array = modified_array * formula_result
    elif isinstance(formula_result, np.ndarray):
        modified_array = modified_array * formula_result
    else:
        raise TypeError(f"公式结果不是数值类型或数组，而是 {type(formula_result)}")

    return modified_array



def calculate_with_boolean_mapping(ecosystem_array, formula_result):
    """
    对生态图数组进行布尔运算，并与公式结果进行计算。

    参数:
    ecosystem_array (numpy.ndarray): 生态图的数组。
    formula_result (Union[int, float, numpy.ndarray]): 公式的结果。

    返回:
    numpy.ndarray: 处理后的计算结果。
    """
    # 对 ecosystem_array 进行布尔运算 ecosystem_array >= 256
    boolean_array = ecosystem_array >= 256

    # 将 formula_result 与布尔数组进行运算（例如，乘法）
    result = boolean_array * formula_result

    return result



# def run_input_matrix_calculator(json_file_path: str, mode: str, formulas: List[Formula] = formulas_list_Q) -> None:
#     """
#     运行矩阵计算器，加载参数并执行指定的公式。
#
#     参数:
#     json_file_path (str): 包含参数的 JSON 文件路径。
#     mode (str): 当前计算模式。
#     formulas (List[Formula]): 可用的公式列表。默认值为 formulas_list_Q。
#
#     返回:
#     None
#
#     异常:
#     Exception: 如果加载文件或执行公式时出错，会打印错误信息。
#     """
#     # 初始化计算器
#     calculator = MatrixCalculator()
#     calculator.load_formulas(formulas)
#     calculator.set_mode(mode)
#
#     # 加载参数
#     parameters = load_parameters(json_file_path)
#     for parameter in parameters:
#         calculator.param_manager.add_parameter(parameter)
#
#     # 获取可用公式
#     available_formulas = calculator.prompt_available_formulas()
#
#     # 提示用户可以选择运行的公式
#     if available_formulas:
#         print("可以运行以下公式：")
#         for formula in available_formulas:
#             print(f"- {formula}")
#     else:
#         print("没有可运行的公式。")
#         return
#
#     # 用户选择运行的公式
#     selected_formula = input("请输入要运行的公式名称：")
#
#     # 运行公式
#     if selected_formula in available_formulas:
#         try:
#             print(f"正在运行公式: {formula}")
#             result = calculator.execute_formula(formula)
#             print(f"未与生态图进行运算的计算结果：\n{result}")
#         except Exception as e:
#             print(f"执行公式 {formula} 时出错: {e}")

def extract_codes(ecosystem):
    """
    递归函数，用于提取生态系统及其子系统的Code。

    参数:
    ecosystem (dict): 生态系统的JSON对象。

    返回:
    list: 包含当前生态系统及其所有子系统Code的列表。
    """
    codes = [ecosystem["code"]]
    for child in ecosystem.get("children", []):
        codes += extract_codes(child)
    return codes

def process_ecosystem_tif(input_tif, output_template=None):
    """
    处理生态系统的TIF文件，按照ecosystems_type.json划分成9大生态系统。

    参数:
    input_tif (str): 输入的TIF文件路径。
    output_template (str): 输出TIF文件的路径模板，包含{name}占位符。
                            默认值为 "./output/ecosystem_{name}.tif"

    返回:
    list: 生成的TIF文件路径列表。
    """
    # 设置默认输出路径模板
    if output_template is None:
        output_template = "./output/ecosystem_{name}.tif"

    # 确保输出目录存在
    output_dir = os.path.dirname(output_template.format(name=""))
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 读取 JSON 数据
    ecosystem_json = load_parameters_from_json(r"F:\code\dev\calc-gep-regulate-cqc\devCalc\data\ecosystems_type.json")

    # 为每个一级生态系统生成Code列表
    ecosystem_codes = {}
    for ecosystem in ecosystem_json["applicable_ecosystems"]:
        name = ecosystem["name"]
        codes = extract_codes(ecosystem)
        ecosystem_codes[name] = codes

    print("ecosystem_codes",ecosystem_codes)

    # 读取原始TIF文件
    with rasterio.open(input_tif) as src:
        data = src.read(1)
        meta = src.meta.copy()
        meta.update(nodata=-9999)  # 设置无效值为-9999

    generated_files = []
    report = []

    # 为每个一级生态系统生成子图
    for name, codes in ecosystem_codes.items():
        # 生成掩膜：属于当前生态系统的区域保留原值，其他设为0
        mask = np.isin(data, codes)
        masked_data = np.where(mask, data, 0).astype(data.dtype)

        # 如果掩膜中没有有效值，跳过生成
        if not np.any(mask):
            report.append(f"未生成文件: {name} (无有效像素)")
            continue

        # 输出路径
        output_path = output_template.format(name=name)

        # 写入新TIF文件
        with rasterio.open(output_path, 'w', **meta) as dst:
            dst.write(masked_data, 1)

        generated_files.append(output_path)
        report.append(f"生成文件: {output_path} (包含Code: {codes})")

    # 打印报告
    print("\n处理报告：")
    for line in report:
        print(line)

    return generated_files

def extract_ecosystem_arrays(input_tif):
    """
    从输入的TIF文件中提取各生态系统类型的numpy数组列表。无某一生态系统则不返回该生态系统。

    参数:
    input_tif (str): 输入的TIF文件路径。

    返回:
    dict: 包含生态系统名称和对应numpy数组的字典。
    """
    # 读取 JSON 数据
    ecosystem_json = load_parameters_from_json(r"F:\code\dev\calc-gep-regulate-cqc\devCalc\data\ecosystems_type.json")

    # 为每个一级生态系统生成Code列表
    ecosystem_codes = {}
    for ecosystem in ecosystem_json["applicable_ecosystems"]:
        name = ecosystem["name"]
        codes = extract_codes(ecosystem)
        ecosystem_codes[name] = codes

    # 读取原始TIF文件
    with rasterio.open(input_tif) as src:
        data = src.read(1)

    ecosystem_arrays = {}

    # 为每个一级生态系统生成子图
    for name, codes in ecosystem_codes.items():
        import time
        start_time = time.time()
        # 生成掩膜：属于当前生态系统的区域保留原值，其他设为255 无效值
        mask = np.isin(data, codes)
        masked_data = np.where(mask, data, np.nan).astype(np.float32) #  NumPy 的类型提升规则  np.nan是float64

        # 将 masked_data 中的 255 替换为 np.nan
        # print("masked_data中无nan", masked_data.dtype)
        # masked_data[masked_data == 255] = np.nan
        print("masked_data有nan", masked_data.dtype)

        end_time = time.time()
        elapsed_time = end_time - start_time
        print(f"处理完: “{name}” 生态图！用时 {elapsed_time:.2f} 秒")

        # 输出最大值和最小值
        max_value = np.nanmax(masked_data)
        min_value = np.nanmin(masked_data)
        print(f"“{name}” 生态图的最大值: {max_value}, 最小值: {min_value}")

        # 如果掩膜中没有有效值，跳过生成
        if not np.any(mask):
            continue

        ecosystem_arrays[name] = masked_data


    return ecosystem_arrays

def get_ecosystem_codes(ecosystem_list, ecosystem_dict):
    """
    根据生态系统名称列表，在字典中获取对应的代码列表。

    参数:
    ecosystem_list (list): 生态系统名称列表。
    ecosystem_dict (dict): 包含生态系统名称和代码的字典。

    返回:
    list: 对应的代码列表。
    """
    result_codes = []
    for ecosystem in ecosystem_list:
        if ecosystem in ecosystem_dict:
            result_codes.extend(ecosystem_dict[ecosystem])
        else:
            print(f"警告：生态系统 '{ecosystem}' 未在字典中找到。")
    return result_codes



def replace_tif_data_with_array(
    tif_path: str,
    array: np.ndarray,
    output_path: str,
    nodata: Optional[Union[int, float]] = -9999
) -> None:
    """
    用一个数组替换TIF图像的数据。

    :param tif_path: 原始TIF文件的路径
    :param array: 替换用的numpy数组
    :param output_path: 替换后保存的TIF文件路径
    :param nodata: 无效值，默认值为-9999
    """
    # 打开TIF文件
    with rasterio.open(tif_path) as src:
        # 获取TIF图像的元信息
        meta = src.meta.copy()
        # 获取TIF图像的数据
        tif_data = src.read()

        # 检查数组大小是否与TIF图像一致
        tif_height, tif_width = src.height, src.width
        array_height, array_width = array.shape

        if tif_height != array_height or tif_width != array_width:
            raise ValueError(f"数组大小（{array_height}x{array_width}）与TIF图像大小（{tif_height}x{tif_width}）不一致！")

        # # 替换TIF图像的数据
        # tif_data[0] = array  # 假设是单波段图像，如果是多波段需要调整

        # 替换TIF图像的数据，并确保使用array的dtype
        tif_data = tif_data.astype(array.dtype)
        tif_data[0] = array  # 假设是单波段图像，如果是多波段需要调整


    # 更新元信息中的数据类型为数组的数据类型
    # todo np.float32
    meta.update(
        dtype=float32,  # rasterio.float32
        nodata=nodata
    )

    # # 设置0为无效值（nodata）
    # # nodata_value = 0
    # meta.update(nodata=nodata)

    # 保存替换后的数据到新的TIF文件
    with rasterio.open(output_path, 'w', **meta) as dst:
        dst.write(tif_data)

    print(f"替换完成，新的TIF文件已保存到 {output_path}")


# 示例用法
# tif_path = r"F:\code\dev\calc-gep-regulate-cqc\devCalc\data\test\test_eco.tif"  # 原始TIF文件路径
# array = np.random.rand(10, 10)  # 创建一个随机数组作为替换数据
# output_path = r"F:\code\dev\calc-gep-regulate-cqc\calc_v1\output\test_replaced_output.tif"  # 输出路径
#
# replace_tif_data_with_array(tif_path, array, output_path)

def find_ecosystems_by_codes(codes):
    # 生态系统字典
    dict_eco = {
        '森林生态系统': [1000, 1100, 1101, 1102, 1200, 1201, 1202, 1300, 1301, 1400, 1401, 1900],
        '灌丛生态系统': [2000, 2100, 2101, 2102, 2200, 2201, 2300, 2301, 2900],
        '草地生态系统': [3000, 3100, 3101, 3102, 3200, 3201, 3202, 3300, 3301, 3302, 3400, 3401, 3402, 3900],
        '湿地生态系统': [4000, 4100, 4101, 4102, 4103, 4200, 4201, 4202, 4203, 4300, 4301, 4302, 4900],
        '农田生态系统-无水': [5000, 5100, 5102, 5200, 5201, 5202, 5900],
        '农田生态系统-水田': [5101],
        '城市生态系统（建成区）': [6000, 6100, 6101, 6102, 6103, 6104, 6200, 6201, 6202, 6203, 6204, 6205, 6206, 6900],
        '城市生态系统-城市水体': [6300, 6301],
        '荒漠生态系统': [7000, 7100, 7101, 7200, 7201, 7300, 7301, 7900],
        '其他陆域': [8000, 8100, 8101, 8200, 8201, 8202, 8900],
        '海洋生态系统': [9000, 9100, 9101, 9102, 9103, 9104, 9105, 9106]
    }
    result = []
    for ecosystem, eco_codes in dict_eco.items():
        if any(code in eco_codes for code in codes):
            result.append(ecosystem)
    return result

